#! /usr/bin/env python
# coding=utf-8

import argparse
import sys
from argparse import RawTextHelpFormatter

from Bio import SeqIO
import csv
import json
import pandas as pd



parser = argparse.ArgumentParser(
    description='''
    读取 mRNA pep 的fasta文件 以及interprotscan结果 和 all_gene.tsv 获得对应的mRNA和pep 产物列表
    用法:
    make_product_table.py -m cves_mRNA.fa -p cves_pep.fa -g all_genes.tsv -i cves_longpep.fa.tsv -a mRNA_for_upload.xlsx -b pep_for_upload.xlsx

    ''',formatter_class=RawTextHelpFormatter)



parser.add_argument('-m',
                help='必须给定，输入的mRNA fasta文件 ')
parser.add_argument('-p',
                help='必须给定，输入的pep fasta文件 ')
parser.add_argument('-g',
                help='必须给定，输入的all gene tsv文件 ')
parser.add_argument('-i',
                help='必须给定，输入的interprotscan tsv 文件 ')
parser.add_argument('-a',
                help='必须给定，输入的mRNA文件 ')
parser.add_argument('-b',
                help='必须给定，输入的pep文件 ')



# 取出pasa中被 merge的 mRNA

args = parser.parse_args()

if not args.i or not args.m or not args.p or not args.g or not args.a or not args.b:
    parser.print_help()
    sys.exit()


# 读取
infile_mRNA = args.m

infile_pep = args.p

infile_gene = args.g

infile_interprot = args.i

outfile_mRNA = args.a

outfile_pep = args.b








inter_prot_dic ={}

with open(infile_interprot) as fila:
	for i in fila:
		k = i.split('#')[0].strip().split('\t')
		if len(k)>5 and k[3] == 'Pfam':
			gene_name = k[0].split('.mRNA')[0]
			if gene_name not in inter_prot_dic:
				inter_prot_dic[gene_name] = {}
			inter_prot_dic[gene_name][k[4]] = [k[6],k[7]]
			

gene_des_dic = {}
with open(infile_gene) as fila:
	head = next(fila).strip().split('\t')
	for i in fila:
		k = i.strip('\n').split('\t')
		if len(k)>2:
			tmp_dic = {x:y for x,y in zip(head,k)}
			gene_des_dic[tmp_dic['gene_name']] = ''
			if tmp_dic['gene_symbol']!='':
				gene_des_dic[tmp_dic['gene_name']] += tmp_dic['gene_symbol']
			
			gene_des_dic[tmp_dic['gene_name']] = gene_des_dic[tmp_dic['gene_name']].strip()

			if gene_des_dic[tmp_dic['gene_name']]!='' and tmp_dic['gene_description'].strip()!='':
				gene_des_dic[tmp_dic['gene_name']] += (' ('+tmp_dic['gene_description']+')')
			elif gene_des_dic[tmp_dic['gene_name']]=='' and tmp_dic['gene_description'].strip()!='':
				gene_des_dic[tmp_dic['gene_name']] += tmp_dic['gene_description']
			
			#gene_des_dic[tmp_dic['gene_name']] =(tmp_dic['gene_symbol'] + ' '+ tmp_dic['gene_description']).strip()

print('finish read gene')



outfile_mRNA_write = pd.ExcelWriter(outfile_mRNA)

df = pd.DataFrame(data=None,columns=['product_id','product_name','type','product_description','gene','product_sequence'])


for i in SeqIO.parse(infile_mRNA,'fasta'):
	mRNA_name = str(i.name)
	gene_name = i.name.split('.mRNA')[0]
	des = ''
	if gene_name in gene_des_dic and gene_des_dic[gene_name]!='':
		des =  gene_des_dic[gene_name]
	

	df = df.append({'product_id':'',
		'product_name':mRNA_name,
		'product_description':des,
		'type':'mRNA',
		'gene':gene_name,
		'product_sequence':str(i.seq)
		},ignore_index=True)

df.to_excel(outfile_mRNA_write, index=False)

outfile_mRNA_write.close()

print('finish write mRNA')



outfile_pep_write = pd.ExcelWriter(outfile_pep)

df = pd.DataFrame(data=None,columns=['product_id','product_name','type','product_description','gene','mRNA','product_sequence','product_pfam'])



for i in SeqIO.parse(infile_pep,'fasta'):
	mRNA_name = str(i.name)
	pep_name = i.name.replace('.mRNA','.pep')
	gene_name = i.name.split('.mRNA')[0]
	pfam = '{}'
	des = ''
	if gene_name in gene_des_dic and gene_des_dic[gene_name]!='':
		des =  gene_des_dic[gene_name]


	if gene_name in inter_prot_dic:

		pfam = json.dumps(inter_prot_dic[gene_name])

	df = df.append({'product_id':'',
	'product_name':pep_name,
	'product_description':des,
	'type':'peptide',
	'gene':gene_name,
	'mRNA':mRNA_name,
	'product_sequence':str(i.seq),
	'product_pfam':pfam,
	},ignore_index=True)

df.to_excel(outfile_pep_write, index=False)

outfile_pep_write.close()

print('finish write pep')
